"""
Author: Morphlng
Date: 2024-04-24 21:51:26
LastEditTime: 2025-02-26 15:08:39
LastEditors: Morphlng
Description: Low-level vehicle control for Carla simulator
FilePath: /DrivingGym/src/driving_gym/environment/agent/action/carla_vehicle_action.py
"""

from __future__ import annotations

from itertools import product

import gymnasium as gym
import numpy as np

from driving_gym.environment.agent.action.action_interface import (
    AbstractAction,
    ActionInterface,
)
from driving_gym.misc.util import override
from driving_gym.simulation.adapter_interface import AdapterInterface


class DirectAction(AbstractAction):
    def __init__(self, action, duration=1):
        super().__init__(action, duration)
        self.action["type"] = "vehicle"

    @override(AbstractAction)
    def run_step(self):
        """Return a carla control signal dict."""
        self.duration -= 1
        return self.action


class CarlaVehicleAction(ActionInterface):
    def __init__(self, action_config: dict, adapter: AdapterInterface):
        """Initialize the action converter for low-level action space

        Args:
            action_config (dict): A dictionary of action config
        """
        super().__init__(action_config, adapter)

        self._is_discrete = self.action_config.get("is_discrete", True)
        if self._is_discrete:
            self._discrete_acc = self.action_config.get(
                "discrete_acc", [-1.0, -0.25, 0, 0.25, 0.5, 0.75]
            )
            self._discrete_steer = self.action_config.get(
                "discrete_steer", [-0.2, 0.0, 0.2]
            )
            self._action_set = {
                idx: action
                for idx, action in enumerate(
                    product(self._discrete_acc, self._discrete_steer)
                )
            }

        self._action_range = self.action_config.get(
            "action_range",
            {
                "throttle": [0, 1.0],
                "brake": [0, 1.0],
                "steer": [-1.0, 1.0],
            },
        )

    @override(ActionInterface)
    def convert_single_action(self, action):
        """Convert the action of a model output to an AbstractAction instance

        Args:
            action: Action input from the model

        Returns:
            DirectAction: A direct action instance
        """
        if self._is_discrete:
            action = self._action_set[action]

        action_dict = {
            "throttle": action[0] if action[0] >= 0 else 0,
            "brake": abs(action[0]) if action[0] < 0 else 0,
            "steer": action[1],
        }
        action_dict = self._clip_action(action_dict)
        return DirectAction(action_dict)

    def _clip_action(self, action_dict: dict):
        """Clip the action to make sure the action is within the range of action space"""
        action_dict["throttle"] = np.clip(
            action_dict["throttle"],
            self._action_range["throttle"][0],
            self._action_range["throttle"][1],
        )
        action_dict["brake"] = np.clip(
            action_dict["brake"],
            self._action_range["brake"][0],
            self._action_range["brake"][1],
        )
        action_dict["steer"] = np.clip(
            action_dict["steer"],
            self._action_range["steer"][0],
            self._action_range["steer"][1],
        )
        return action_dict

    @override(ActionInterface)
    def get_action_space(self):
        if not self._is_discrete:
            return gym.spaces.Box(
                low=np.array(
                    [-self._action_range["brake"][1], self._action_range["steer"][0]]
                ),
                high=np.array(
                    [self._action_range["throttle"][1], self._action_range["steer"][1]]
                ),
                dtype=np.float32,
            )
        else:
            return gym.spaces.Discrete(len(self._action_set))

    @override(ActionInterface)
    def get_mask_space(self):
        return None

    @staticmethod
    def sort_action_set(action_set: dict):
        """Re-arrange the action_set to make sure the order of actions is consistent

        This function can help to make sure that order of low-level action_space is ordered by "throttle/brake" (from low to high) and then "steer" (from low to high)

        By using this, we make sure the brake action is always the first action (index 0)
        """
        sorted_items = sorted(action_set.items(), key=lambda item: item[1][0])
        return {new_key: value for new_key, (old_key, value) in enumerate(sorted_items)}
